# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""TREC2014 Dataset"""
from pathlib import Path
import numpy as np
import argparse
import os
from scipy.sparse import csr_matrix

from mindspore_gl.graph import MindHomoGraph, CsrAdj
from src.base_dataset import BaseDataSet
from src.dataset import Dataset


class TREC(BaseDataSet):
    def __init__(self, root, dataset):
        self._root = Path(root)

        self._csr_row = None
        self._csr_col = None
        self._nodes = None

        self._node_feat = None
        self._node_label = None

        self._train_mask = None
        self._val_mask = None
        self._test_mask = None
        self._npz_file = None

        self.dataset = dataset

        if self._root.is_dir():
            self._preprocess()
        else:
            raise Exception('data file does not exist')

    def _preprocess(self):
        """process data"""
        # qid
        qid_index = np.load('./data/dgat_qid_edge_index.npy')

        row_qid = qid_index[0][:]
        col_qid = qid_index[1][:]
        self.qid_row = len(row_qid)
        self.col_qid = col_qid

        qid = np.ones(row_qid.shape)
        qid_node_count = self.dataset.query_size
        qid_csr_mat = csr_matrix((qid, (row_qid, col_qid)), shape=(qid_node_count + 1, qid_node_count + 1))
        self.qid_csr_mat = qid_csr_mat
        qid_indices = qid_csr_mat.indices
        qid_indpter = qid_csr_mat.indptr

        np.savez(self._root, edge_array=qid_index, adj_csr_indptr=qid_indpter, adj_csr_indices=qid_indices,
                 node_count=qid_node_count, train_mask=self.train_mask,
                 test_mask=self.test_mask, val_mask=self.val_mask)

        # uid
        uid_index = np.load('./data/dgat_uid_edge_index.npy')

        row_uid = uid_index[0][:]
        col_uid = uid_index[1][:]
        uid = np.ones(row_uid.shape)
        self.row_uid = row_uid
        self.uid_row = len(row_uid)

        uid_node_count = self.dataset.doc_size

        uid_csr_mat = csr_matrix((uid, (row_uid, col_uid)), shape=(uid_node_count + 1, uid_node_count + 1))
        self.uid_csr_mat = uid_csr_mat
        uid_indices = uid_csr_mat.indices
        uid_indpter = uid_csr_mat.indptr

        np.savez(self._root, edge_array=uid_index, adj_csr_indptr=uid_indpter, adj_csr_indices=uid_indices,
                 node_count=uid_node_count, train_mask=self.train_mask,
                 test_mask=self.test_mask, val_mask=self.val_mask)
    
    @property
    def train_mask(self):
        """
        Mask of training nodes.

        Returns:
            - numpy.ndarray, array of mask.

        Examples:
            >>> #dataset is an instance object of Dataset
            >>> train_mask = dataset.train_mask
        """
        if self._train_mask is None:
            return self._train_mask

    @property
    def test_mask(self):
        """
        Mask of training nodes.

        Returns:
            - numpy.ndarray, array of mask.

        Examples:
            >>> #dataset is an instance object of Dataset
            >>> test_mask = dataset.test_mask
        """
        if self._test_mask is None:
            return self._test_mask

    @property
    def val_mask(self):
        """
        Mask of training nodes.

        Returns:
            - numpy.ndarray, array of mask.

        Examples:
            >>> #dataset is an instance object of Dataset
            >>> val_mask = dataset.val_mask
        """
        if self._val_mask is None:
            return self._val_mask

    @property
    def qid_node_count(self):
        """
        Number of qid nodes.

        Returns:
            - int, length of csr row.

        Examples:
            >>> #dataset is an instance object of Dataset
            >>> node_count = dataset.node_count
        """
        return self.dataset.query_size

    @property
    def uid_node_count(self):
        """
        Number of uid nodes.

        Returns:
            - int, length of csr row.

        Examples:
            >>> #dataset is an instance object of Dataset
            >>> node_count = dataset.node_count
        """
        return self.dataset.doc_size

    @property
    def qid_edge_count(self):
        """
        Number of qid edges.

        Returns:
            - int, length of csr col.

        Examples:
            >>> #dataset is an instance object of Dataset
            >>> edge_count = dataset.edge_count
        """
        return self.qid_row

    @property
    def uid_edge_count(self):
        """
        Number of uid edges.

        Returns:
            - int, length of csr col.

        Examples:
            >>> #dataset is an instance object of Dataset
            >>> edge_count = dataset.edge_count
        """
        return self.uid_row

    def qid_(self):
        qid_graph = MindHomoGraph()
        qid_node_dict = {idx: idx for idx in range(self.qid_node_count)}
        qid_edge_count = self.col_qid.shape[0]
        qid_edge_ids = np.array(list(range(qid_edge_count))).astype(np.int32)
        qid_graph.set_topo(CsrAdj(self.qid_csr_mat.indptr.astype(np.int32), self.qid_csr_mat.indices.astype(np.int32)),
                           qid_node_dict, qid_edge_ids)
        return qid_graph
    
    def uid_(self):
        uid_graph = MindHomoGraph()
        uid_node_dict = {idx: idx for idx in range(self.uid_node_count)}
        uid_edge_count = self.row_uid.shape[0]
        uid_edge_ids = np.array(list(range(uid_edge_count))).astype(np.int32)
        uid_graph.set_topo(CsrAdj(self.uid_csr_mat.indptr.astype(np.int32), self.uid_csr_mat.indices.astype(np.int32)),
                           uid_node_dict, uid_edge_ids)
        return uid_graph
    
    def data1_(self):
        return self.dataset
    
    def __getitem__(self, idx):
        qid_graph = MindHomoGraph()
        qid_node_dict = {idx: idx for idx in range(self.qid_node_count)}
        qid_edge_count = self.col_qid.shape[0]
        qid_edge_ids = np.array(list(range(qid_edge_count))).astype(np.int32)
        qid_graph.set_topo(CsrAdj(self.qid_csr_mat.indptr.astype(np.int32), self.qid_csr_mat.indices.astype(np.int32)),
                           qid_node_dict, qid_edge_ids)

        uid_graph = MindHomoGraph()
        uid_node_dict = {idx: idx for idx in range(self.uid_node_count)}
        uid_edge_count = self.row_uid.shape[0]
        uid_edge_ids = np.array(list(range(uid_edge_count))).astype(np.int32)
        uid_graph.set_topo(CsrAdj(self.uid_csr_mat.indptr.astype(np.int32), self.uid_csr_mat.indices.astype(np.int32)),
                           uid_node_dict, uid_edge_ids)
        return qid_graph, uid_graph, self.dataset


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='TREC2014')
    parser.add_argument('--max_d_num', type=int, default=10,
                        help='max number of docs in a session')

    path_settings = parser.add_argument_group('path settings')
    path_settings.add_argument('--gpu_num', type=int, default=1,
                               help='gpu_num')
    path_settings.add_argument('--dataset', default='TREC2014',
                               help='name of the dataset to be used')

    args = parser.parse_args()
    dataset = Dataset(args)
    data1, data2, data3 = TREC(root='../data', dataset=dataset)
    print(data3.query_size)
    data._preprocess()